"""
内存管理器
=========

提供内存优化和管理功能。
"""

import numpy as np
from typing import Dict, Tuple, Optional, Any
import gc
import psutil
import os


class MemoryManager:
    """
    内存管理器类
    
    负责内存分配、重用和监控。
    """
    
    def __init__(self, max_memory_mb: float = 1024.0):
        """
        初始化内存管理器
        
        Parameters:
        -----------
        max_memory_mb : float
            最大内存使用量（MB）
        """
        self.max_memory_mb = max_memory_mb
        self.buffer_pool = {}
        self.allocation_stats = {
            'total_allocations': 0,
            'total_deallocations': 0,
            'peak_memory_mb': 0.0,
            'current_memory_mb': 0.0
        }
    
    def allocate_buffers(self, shape: Tuple[int, ...], 
                        dtype: np.dtype = np.float64,
                        count: int = 1) -> np.ndarray:
        """
        分配缓冲区
        
        Parameters:
        -----------
        shape : Tuple[int, ...]
            数组形状
        dtype : np.dtype
            数据类型
        count : int
            分配数量
            
        Returns:
        --------
        np.ndarray
            分配的缓冲区
        """
        buffer_key = (shape, dtype)
        
        if buffer_key not in self.buffer_pool:
            self.buffer_pool[buffer_key] = []
        
        pool = self.buffer_pool[buffer_key]
        
        if len(pool) >= count:
            # 从池中获取
            buffers = [pool.pop() for _ in range(count)]
        else:
            # 创建新的缓冲区
            needed = count - len(pool)
            existing = [pool.pop() for _ in range(len(pool))]
            new_buffers = [np.empty(shape, dtype=dtype) for _ in range(needed)]
            buffers = existing + new_buffers
            
            self.allocation_stats['total_allocations'] += needed
        
        if count == 1:
            return buffers[0]
        else:
            return buffers
    
    def reuse_arrays(self, arrays: list):
        """
        回收数组到缓冲池
        
        Parameters:
        -----------
        arrays : list
            要回收的数组列表
        """
        if not isinstance(arrays, list):
            arrays = [arrays]
        
        for arr in arrays:
            if isinstance(arr, np.ndarray):
                buffer_key = (arr.shape, arr.dtype)
                
                if buffer_key not in self.buffer_pool:
                    self.buffer_pool[buffer_key] = []
                
                # 清零数组（可选）
                arr.fill(0)
                
                self.buffer_pool[buffer_key].append(arr)
                self.allocation_stats['total_deallocations'] += 1
    
    def garbage_collect(self):
        """执行垃圾回收"""
        collected = gc.collect()
        return collected
    
    def monitor_memory_usage(self) -> Dict[str, float]:
        """
        监控内存使用情况
        
        Returns:
        --------
        Dict[str, float]
            内存使用信息
        """
        try:
            process = psutil.Process(os.getpid())
            memory_info = process.memory_info()
            
            current_memory_mb = memory_info.rss / (1024 * 1024)
            
            # 更新统计
            self.allocation_stats['current_memory_mb'] = current_memory_mb
            if current_memory_mb > self.allocation_stats['peak_memory_mb']:
                self.allocation_stats['peak_memory_mb'] = current_memory_mb
            
            return {
                'current_memory_mb': current_memory_mb,
                'peak_memory_mb': self.allocation_stats['peak_memory_mb'],
                'max_memory_mb': self.max_memory_mb,
                'memory_usage_ratio': current_memory_mb / self.max_memory_mb,
                'available_memory_mb': self.max_memory_mb - current_memory_mb
            }
            
        except Exception as e:
            return {
                'error': str(e),
                'current_memory_mb': 0.0,
                'peak_memory_mb': 0.0,
                'max_memory_mb': self.max_memory_mb,
                'memory_usage_ratio': 0.0,
                'available_memory_mb': self.max_memory_mb
            }
    
    def check_memory_limit(self) -> bool:
        """
        检查是否超过内存限制
        
        Returns:
        --------
        bool
            是否超过限制
        """
        memory_info = self.monitor_memory_usage()
        return memory_info['memory_usage_ratio'] > 0.9  # 90%阈值
    
    def optimize_array_layout(self, arrays: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
        """
        优化数组内存布局
        
        Parameters:
        -----------
        arrays : Dict[str, np.ndarray]
            输入数组字典
            
        Returns:
        --------
        Dict[str, np.ndarray]
            优化后的数组字典
        """
        optimized = {}
        
        for key, arr in arrays.items():
            if isinstance(arr, np.ndarray):
                # 确保C连续布局
                if not arr.flags['C_CONTIGUOUS']:
                    optimized[key] = np.ascontiguousarray(arr)
                else:
                    optimized[key] = arr
                
                # 优化数据类型（如果合适）
                if arr.dtype == np.float64 and np.all(np.abs(arr) < 1e6):
                    # 可以考虑使用float32来节省内存
                    # optimized[key] = optimized[key].astype(np.float32)
                    pass
            else:
                optimized[key] = arr
        
        return optimized
    
    def get_buffer_pool_info(self) -> Dict[str, Any]:
        """
        获取缓冲池信息
        
        Returns:
        --------
        Dict[str, Any]
            缓冲池信息
        """
        pool_info = {}
        total_buffers = 0
        total_memory_mb = 0.0
        
        for (shape, dtype), buffers in self.buffer_pool.items():
            buffer_count = len(buffers)
            buffer_size_mb = (np.prod(shape) * dtype.itemsize * buffer_count) / (1024 * 1024)
            
            pool_info[f"{shape}_{dtype}"] = {
                'count': buffer_count,
                'size_mb': buffer_size_mb,
                'shape': shape,
                'dtype': str(dtype)
            }
            
            total_buffers += buffer_count
            total_memory_mb += buffer_size_mb
        
        return {
            'pool_details': pool_info,
            'total_buffers': total_buffers,
            'total_memory_mb': total_memory_mb,
            'pool_types': len(self.buffer_pool)
        }
    
    def clear_buffer_pool(self):
        """清空缓冲池"""
        self.buffer_pool.clear()
        self.garbage_collect()
    
    def get_allocation_stats(self) -> Dict[str, Any]:
        """
        获取分配统计信息
        
        Returns:
        --------
        Dict[str, Any]
            分配统计
        """
        stats = self.allocation_stats.copy()
        stats.update(self.monitor_memory_usage())
        stats.update(self.get_buffer_pool_info())
        
        return stats
    
    def suggest_memory_optimization(self) -> Dict[str, Any]:
        """
        建议内存优化方案
        
        Returns:
        --------
        Dict[str, Any]
            优化建议
        """
        memory_info = self.monitor_memory_usage()
        buffer_info = self.get_buffer_pool_info()
        
        suggestions = []
        
        # 内存使用率过高
        if memory_info['memory_usage_ratio'] > 0.8:
            suggestions.append("内存使用率过高，建议增加内存限制或减少数据规模")
        
        # 缓冲池过大
        if buffer_info['total_memory_mb'] > 100:
            suggestions.append("缓冲池占用内存过多，建议定期清理")
        
        # 分配/释放不平衡
        alloc_ratio = (self.allocation_stats['total_deallocations'] / 
                      max(self.allocation_stats['total_allocations'], 1))
        if alloc_ratio < 0.5:
            suggestions.append("内存分配和释放不平衡，可能存在内存泄漏")
        
        return {
            'memory_info': memory_info,
            'buffer_info': buffer_info,
            'suggestions': suggestions,
            'optimization_score': self._calculate_optimization_score(memory_info, buffer_info)
        }
    
    def _calculate_optimization_score(self, memory_info: Dict[str, float], 
                                    buffer_info: Dict[str, Any]) -> float:
        """计算优化分数（0-100）"""
        score = 100.0
        
        # 内存使用率惩罚
        if memory_info['memory_usage_ratio'] > 0.8:
            score -= (memory_info['memory_usage_ratio'] - 0.8) * 200
        
        # 缓冲池效率奖励
        if buffer_info['total_buffers'] > 0:
            score += min(buffer_info['total_buffers'] * 2, 20)
        
        return max(0.0, min(100.0, score))


# 全局内存管理器实例
_global_memory_manager = None


def get_memory_manager() -> MemoryManager:
    """获取全局内存管理器实例"""
    global _global_memory_manager
    if _global_memory_manager is None:
        _global_memory_manager = MemoryManager()
    return _global_memory_manager


def set_memory_limit(max_memory_mb: float):
    """设置内存限制"""
    manager = get_memory_manager()
    manager.max_memory_mb = max_memory_mb


def monitor_memory() -> Dict[str, float]:
    """监控内存使用"""
    manager = get_memory_manager()
    return manager.monitor_memory_usage()


def optimize_arrays(arrays: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
    """优化数组内存布局"""
    manager = get_memory_manager()
    return manager.optimize_array_layout(arrays)